package zjl.graph;

import java.util.*;

public class DijkstraShortestPathAlgorithm {

    // input
    private BaseGraph graph;

    // intermediate variables
    private Set<BaseVertex> visitedVertexSet = new HashSet<>();
    private Queue<BaseVertex> vertexCandidateQueue = new PriorityQueue<>();
    private Map<BaseVertex, Double> startVertexDistanceIndex = new HashMap<>();
    private Map<BaseVertex, BaseVertex> predecessorIndex = new HashMap<>();

    public DijkstraShortestPathAlgorithm(BaseGraph graph) {
        this.graph = graph;
    }

    public void clear() {
        this.visitedVertexSet.clear();
        this.vertexCandidateQueue.clear();
        this.startVertexDistanceIndex.clear();
        this.predecessorIndex.clear();
    }

    public Map<BaseVertex, Double> getStartVertexDistanceIndex() {
        return this.startVertexDistanceIndex;
    }

    public Map<BaseVertex, BaseVertex> getPredecessorIndex() {
        return this.predecessorIndex;
    }

    // Construct a tree rooted at "root" with the shortest paths to the other vertices
    public void getShortestPathTree(BaseVertex root) {
        this.determineShortestPaths(root, null, true);
    }

    // Construct a flower rooted at "root" with the shortest paths from the other vertices
    public void getShortestPathFlower(BaseVertex root) {
        this.determineShortestPaths(null, root, false);
    }

    protected void determineShortestPaths(BaseVertex sourceVertex, BaseVertex sinkVertex, boolean isSource2Sink) {
        // 0. clear up variables
        clear();

        // 1. initialize members
        BaseVertex endVertex = isSource2Sink ? sinkVertex : sourceVertex;
        BaseVertex startVertex = isSource2Sink ? sourceVertex : sinkVertex;
        this.startVertexDistanceIndex.put(startVertex, 0.0);
        startVertex.setWeight(0.0);

        this.vertexCandidateQueue.offer(startVertex);

        // 2. start searching for the shortest path
        while (!vertexCandidateQueue.isEmpty()) {
            BaseVertex currCandidate = this.vertexCandidateQueue.poll();

            if (currCandidate.equals(endVertex)) {
                break;
            }

            visitedVertexSet.add(currCandidate);

            updateVertex(currCandidate, isSource2Sink);
        }

    }

    private void updateVertex(BaseVertex vertex, boolean isSource2Sink) {
        // 1. get the neighboring vertices
        Set<BaseVertex> neighborVertexList = isSource2Sink ? graph.getAdjacentVertices(vertex)
                                                            : graph.getPrecedentVertices(vertex);

        // 2. update the distance passing on current vertex
        for (BaseVertex currAdjacentVertex : neighborVertexList) {
            // 2.1 skip if visited before
            if (this.visitedVertexSet.contains(currAdjacentVertex)) {
                continue;
            }

            // 2.2 calculate the new distance
            double distance = startVertexDistanceIndex.containsKey(vertex) ? startVertexDistanceIndex.get(vertex)
                                                                           : Graph.DISCONNECTED;

            distance += isSource2Sink ? graph.getEdgeWeight(vertex, currAdjacentVertex)
                                      : graph.getEdgeWeight(currAdjacentVertex, vertex);

            // 2.3 update the distance if necessary
            if (!startVertexDistanceIndex.containsKey(currAdjacentVertex)
                || startVertexDistanceIndex.get(currAdjacentVertex) > distance) {

                startVertexDistanceIndex.put(currAdjacentVertex, distance);

                predecessorIndex.put(currAdjacentVertex, vertex);
                currAdjacentVertex.setWeight(distance);
                vertexCandidateQueue.add(currAdjacentVertex);
            }
        }
    }

    public Path getShortestPath(BaseVertex sourceVertex, BaseVertex sinkVertex) {
        determineShortestPaths(sourceVertex, sinkVertex, true);

        List<BaseVertex> vertexList = new ArrayList<>();
        List<BaseEdge> edgeList = new ArrayList<>();

        double weight = startVertexDistanceIndex.containsKey(sinkVertex) ? startVertexDistanceIndex.get(sinkVertex)
                                                                         : Graph.DISCONNECTED;
        if (weight != Graph.DISCONNECTED) {
            BaseVertex postVertex = null;
            BaseVertex currVertex = sinkVertex;
            do {
                vertexList.add(currVertex);
                postVertex = currVertex;
                currVertex = predecessorIndex.get(currVertex);
                edgeList.add(this.graph.getEdgeInstance(currVertex, postVertex));
            } while (currVertex != null && currVertex != sourceVertex);

            vertexList.add(sourceVertex);

            Collections.reverse(vertexList);
            Collections.reverse(edgeList);
        }

        return new Path(vertexList, edgeList, weight);
    }

    // calculate the distance from the target vertex to the input vertex using forward star form.
    public Path updateCostForward(BaseVertex vertex) {
        double cost = Graph.DISCONNECTED;

        // 1. get the set of successors of the input vertex
        Set<BaseVertex> adjVertexSet = graph.getAdjacentVertices(vertex);

        // 2. make sure the input vertex exists in the index
        if (!this.startVertexDistanceIndex.containsKey(vertex)) {
            this.startVertexDistanceIndex.put(vertex, Graph.DISCONNECTED);
        }

        // 3. update the distance from the root to the input vertex if necessary
        for (BaseVertex currVertex : adjVertexSet) {
            // 3.1 get the distance from the root to one successor of the input vertex
            double distance = this.startVertexDistanceIndex.containsKey(currVertex) ? startVertexDistanceIndex.get(currVertex)
                                                                                    : Graph.DISCONNECTED;

            // 3.2 calculate the distance from the root to the input vertex
            distance += this.graph.getEdgeWeight(vertex, currVertex);

            // 3.3 update the distance if necessary
            double costOfVertex = this.startVertexDistanceIndex.get(vertex);
            if (costOfVertex > distance) {
                this.startVertexDistanceIndex.put(vertex, distance);
                this.predecessorIndex.put(vertex, currVertex);
                cost = distance;
            }
        }

        // 4. create the subpath if exists
        Path subPath = null;
        if (cost < Graph.DISCONNECTED) {
            subPath = new Path();
            subPath.setWeight(cost);

            List<BaseVertex> vertexList = subPath.getVertexList();
            List<BaseEdge> edgeList = subPath.getEdgeList();

            BaseVertex currVertex = vertex;
            BaseVertex selVertex = this.predecessorIndex.get(vertex);

            vertexList.add(currVertex);
            edgeList.add(this.graph.getEdgeInstance(currVertex, selVertex));

            while (selVertex != null) {
                currVertex = selVertex;
                selVertex = this.predecessorIndex.get(selVertex);

                vertexList.add(currVertex);
                if (selVertex != null) {
                    edgeList.add(this.graph.getEdgeInstance(currVertex, selVertex));
                }
            }
        }

        return subPath;
    }

    // correct costs of successors of the input vertex using backward star form.
    public void correctCostBackward(BaseVertex vertex) {
        // 1. initialize the list of vertex to be updated
        List<BaseVertex> vertexList = new LinkedList<>();
        vertexList.add(vertex);

        // 2. update the cost of relevant precedents of the input vertex
        while (!vertexList.isEmpty()) {
            BaseVertex currVertex = vertexList.remove(0);
            double costOfCurrVertex = this.startVertexDistanceIndex.get(currVertex);

            Set<BaseVertex> preVertexSet = this.graph.getPrecedentVertices(currVertex);
            for (BaseVertex preVertex : preVertexSet) {
                double costOfPreVertex = this.startVertexDistanceIndex.containsKey(preVertex) ? startVertexDistanceIndex.get(preVertex)
                                                                                              : Graph.DISCONNECTED;

                double freshCost = costOfCurrVertex + this.graph.getEdgeWeight(preVertex, currVertex);
                if (costOfPreVertex > freshCost) {
                    this.startVertexDistanceIndex.put(preVertex, freshCost);
                    this.predecessorIndex.put(preVertex, currVertex);
                    vertexList.add(preVertex);
                }
            }
        }
    }
}
